import svar
import cv2
import numpy as np
import time
import os
import math

height=0

def callback_image(image):
  #print(image)
  m= np.array(image)
  m.dtype='uint8'
  print(m.shape,m.dtype.name) # BGRA

  cv2.imshow('video', m)
  cv2.waitKey(10)

rtmv=open(str(time.time())+".rtmv",'wb',buffering=0)
#h264=open(str(time.time())+".h264",'wb',buffering=0)

def callback_buf(buf):
  #pass
  m= np.array(buf)
  mem=memoryview(buf)
  #print(m.shape,m.dtype.name,m.shape[0],mem) # BGRA
  #if m.shape[0] > 128:
    #print('h264',m.shape[0])
    #h264.write(mem.tobytes())
  if height > 100:
    rtmv.write(mem.tobytes())# only record height above 100m

osdk=svar.load('svar_osdk')

env=osdk.LinuxSetup(True)
vehicle = env.getVehicle()

options={}
options["camera_position"]=osdk.OSDK_CAMERA_POSITION_NO_1
options["playload_index"]=osdk.PAYLOAD_INDEX_0
options["h264_source"]=osdk.OSDK_CAMERA_SOURCE_H20T_ZOOM
options["enable_rtmv"]=True
node=osdk.VehicleNode(vehicle,options)
sub_rtmv=osdk.messenger.subscribe("rtmv",0,callback_buf)
broadcast=vehicle.broadcast
gimbalManager=vehicle.gimbalManager

def rad2degree(r):
  return 180/3.1415926*r

def degree2rad(d):
  return 3.1415926/180.*d

f_w=20.
f_h=15.
h=300
v_max=15.

o_f=0.6
o_s=0.6
o_i=0.6
n=4

w_i=2*h*math.tan(degree2rad(f_w)/2)
h_i=2*h*math.tan(degree2rad(f_h)/2)
f_h_e=f_h*o_i+(1-o_i)*f_h*n
t=(1-o_f)*w_i*2/v_max
d_s=2*h*math.tan(degree2rad(f_h_e)/2)- o_s*h_i
tm=t/2/n
roll= rad2degree(math.atan((n-1)/n*w_i/2/h))

print("w_i:%f,h_i:%f,fhe:%f,t:%f,d_s:%f,roll:%f"%(w_i,h_i,f_h_e,t,d_s,roll))

#compute control points t-pitch-roll

num=10*n
te1=(t/2-tm)/2
te2=te1+t/2
v_pitch= (1-o_i)*f_h/(te1*2)*(n-1)
v_roll = -2*math.atan(te1*v_max/h)/tm
times=[]
pitchs=[]
rolls=[]

while True:
  q=broadcast.getQuaternion()
  global_pos=broadcast.getGlobalPosition()
  height=global_pos["height"]
  uav_rotation=osdk.SO3(q[0],q[1],q[2],q[3])
  #print("x:",uav_rotation.trans(osdk.Point3d(1,0,0)))
  #print("y:",uav_rotation.trans(osdk.Point3d(0,1,0)))
  #print("z:",uav_rotation.trans(osdk.Point3d(0,0,1)))
  x_dir=uav_rotation.trans(osdk.Point3d(1,0,0))
  uav_yaw=rad2degree(math.atan2(x_dir.y,x_dir.x))
  print("uav_yaw",uav_yaw,"height",height)
  t_i= time.time()
  t_i= t_i%t
  if t_i <= t/2-tm:
    droll = math.atan((t_i-te1)*v_max/h)
    dpitch= v_pitch*(t_i-te1)
  elif t_i<=t/2:
    droll = v_roll*(t_i-t/2+tm/2)
    dpitch= v_pitch*(t/2-tm-te1)
  elif t_i <= t-tm:
    droll = math.atan((t_i-te2)*v_max/h)
    dpitch= -v_pitch*(t_i-te2)
  else:
    droll = v_roll*(t_i-t+tm/2)
    dpitch= -v_pitch*(t-tm-te2)
  pitch=dpitch-90
  roll=rad2degree(droll)
  print("%f, %f, %f"%(t_i,pitch,roll))
  print(broadcast.getGimbal())
  gimbalManager.rotateSync(osdk.PAYLOAD_INDEX_0,{"pitch":pitch,"yaw":uav_yaw-90,"roll":roll,"time":0.1},1)
  time.sleep(0.1)